#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
屏幕覆盖层模块
实现在屏幕上直接显示检测结果，不使用额外窗口
"""

import cv2
import numpy as np
import time
import os
import threading
import queue
import pyautogui
import tkinter as tk
from PIL import Image, ImageTk, ImageDraw, ImageFont
import platform
import sys

class TransparentOverlay:
    """创建一个透明覆盖层窗口，用于直接在屏幕上显示检测结果"""
    
    def __init__(self, alpha=0.6, line_thickness=3, font_scale=1.0, font_thickness=2):
        """
        初始化覆盖层
        
        参数:
            alpha: 覆盖层透明度 (0-1)，增加到0.6让内容更明显
            line_thickness: 边界框线条粗细，增加到3
            font_scale: 字体大小缩放
            font_thickness: 字体粗细
        """
        self.alpha = alpha
        self.line_thickness = line_thickness
        self.font_scale = font_scale
        self.font_thickness = font_thickness
        self.fps = 0
        self.running = False
        self.frame_queue = queue.Queue(maxsize=1)  # 只保留最新的一帧
        self.root = None
        self.canvas = None
        self.tk_image = None
        self.screen_width, self.screen_height = pyautogui.size()
        
        # 检测操作系统
        self.system = platform.system().lower()
        print(f"检测到操作系统: {self.system}")
        
        # 尝试加载中文字体
        self.font = self._load_appropriate_font()
        
        # 定义颜色映射 (BGR格式)
        self.colors = {
            'red': (0, 0, 255),
            'green': (0, 255, 0), 
            'blue': (255, 0, 0),
            'yellow': (0, 255, 255),
            'cyan': (255, 255, 0),
            'magenta': (255, 0, 255),
            'white': (255, 255, 255),
            'black': (0, 0, 0)
        }
    
    def _load_appropriate_font(self):
        """根据系统加载合适的字体"""
        # 尝试常见字体路径
        font_paths = []
        
        if self.system == 'windows':
            # Windows 字体路径
            font_paths = [
                os.path.join(os.environ.get('WINDIR', 'C:\\Windows'), 'Fonts\\simhei.ttf'),
                os.path.join(os.environ.get('WINDIR', 'C:\\Windows'), 'Fonts\\msyh.ttc'),
                os.path.join(os.environ.get('WINDIR', 'C:\\Windows'), 'Fonts\\arial.ttf')
            ]
        elif self.system == 'linux':
            # Linux 字体路径
            font_paths = [
                '/usr/share/fonts/truetype/droid/DroidSansFallbackFull.ttf',
                '/usr/share/fonts/noto/NotoSansCJK-Regular.ttc',
                '/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf'
            ]
        elif self.system == 'darwin':
            # MacOS 字体路径
            font_paths = [
                '/System/Library/Fonts/PingFang.ttc',
                '/Library/Fonts/Arial Unicode.ttf',
                '/Library/Fonts/Arial.ttf'
            ]
        
        # 尝试逐个加载字体
        for font_path in font_paths:
            try:
                if os.path.exists(font_path):
                    return ImageFont.truetype(font_path, int(16 * self.font_scale))
            except Exception as e:
                print(f"无法加载字体 {font_path}: {e}")
        
        # 如果都失败，使用默认字体
        try:
            return ImageFont.load_default()
        except Exception as e:
            print(f"无法加载默认字体: {e}")
            return None
    
    def start(self):
        """启动覆盖层窗口"""
        self.running = True
        
        # 创建并配置Tkinter窗口
        self.root = tk.Tk()
        self.root.title("YOLOv8检测覆盖层")
        
        # 根据不同操作系统设置不同的窗口属性
        try:
            if self.system == 'windows':
                # Windows系统
                self.root.attributes('-alpha', self.alpha)
                self.root.attributes('-topmost', True)
                self.root.attributes('-transparentcolor', 'black')
                self.root.overrideredirect(True)  # 无边框窗口
            elif self.system == 'darwin':
                # MacOS系统
                self.root.attributes('-alpha', self.alpha)
                self.root.attributes('-topmost', True)
                self.root.overrideredirect(True)
            else:
                # Linux及其他系统
                self.root.attributes('-alpha', self.alpha)
                self.root.attributes('-topmost', True)
                self.root.overrideredirect(True)
            
            print("✓ 屏幕覆盖层窗口配置成功")
        except Exception as e:
            print(f"警告: 无法设置窗口属性: {e}")
            print("尝试使用备用设置...")
            # 备用设置
            try:
                self.root.attributes('-alpha', self.alpha)
                self.root.attributes('-topmost', True)
            except:
                print("警告: 无法设置基本窗口属性，覆盖层可能无法正常显示")
        
        # 设置全屏
        self.root.geometry(f"{self.screen_width}x{self.screen_height}+0+0")
        self.root.configure(background='black')
        
        # 创建快捷键
        self.root.bind('<Escape>', lambda e: self.stop())  # ESC键退出
        self.root.bind('q', lambda e: self.stop())  # q键退出
        self.root.bind('+', lambda e: self.adjust_alpha(0.05))  # +键增加透明度
        self.root.bind('-', lambda e: self.adjust_alpha(-0.05))  # -键减少透明度
        
        # 创建Canvas用于显示检测结果
        self.canvas = tk.Canvas(self.root, width=self.screen_width, height=self.screen_height, 
                               highlightthickness=0, bg='black')
        self.canvas.pack()
        
        # 启动更新线程
        threading.Thread(target=self._update_loop, daemon=True).start()
        
        # 显示初始提示
        self.show_help()
        
        # 启动Tkinter主循环
        try:
            self.root.mainloop()
        except Exception as e:
            print(f"窗口主循环错误: {e}")
            self.stop()
    
    def stop(self):
        """停止覆盖层窗口"""
        self.running = False
        if self.root:
            try:
                self.root.quit()
                self.root.destroy()
            except RuntimeError:
                # 如果在非主线程调用，忽略错误
                pass
    
    def adjust_alpha(self, delta):
        """调整透明度"""
        self.alpha = max(0.1, min(1.0, self.alpha + delta))
        if self.root:
            self.root.attributes('-alpha', self.alpha)
    
    def update_frame(self, detections, fps):
        """更新检测结果帧"""
        if not self.running:
            return
        
        self.fps = fps
        
        # 确保队列中只保留最新帧
        while not self.frame_queue.empty():
            try:
                self.frame_queue.get_nowait()
            except queue.Empty:
                break
        
        # 将新帧放入队列
        self.frame_queue.put(detections)
    
    def _update_loop(self):
        """更新显示的线程函数"""
        last_update = time.time()
        last_detections = None
        
        while self.running:
            # 限制更新频率，避免闪烁
            now = time.time()
            if now - last_update < 0.1:  # 降低更新频率到10fps，减少闪烁
                time.sleep(0.05)
                continue
            
            last_update = now
            
            try:
                if not self.frame_queue.empty():
                    detections = self.frame_queue.get()
                    # 如果检测结果没有变化，不更新画面以减少闪烁
                    if last_detections != detections:
                        self._render_detections(detections)
                        last_detections = detections
            except Exception as e:
                print(f"更新覆盖层时出错: {e}")
    
    def _render_detections(self, detections):
        """渲染检测结果到覆盖层"""
        if not self.running or not self.canvas:
            return
        
        try:
            # 创建透明图像
            overlay = Image.new('RGBA', (self.screen_width, self.screen_height), (0, 0, 0, 0))
            draw = ImageDraw.Draw(overlay)
            
            # 绘制检测结果
            if detections:
                # 为不同类别选择不同颜色
                class_colors = {}
                color_list = list(self.colors.values())
                
                for det in detections:
                    # 获取类别、置信度和边界框
                    label = det['class']
                    confidence = det['confidence']
                    x1, y1, x2, y2 = det['box']
                    
                    # 为类别分配颜色
                    if label not in class_colors:
                        color_idx = len(class_colors) % len(color_list)
                        class_colors[label] = color_list[color_idx]
                    
                    color = class_colors[label]
                    
                    # BGR转RGB并添加Alpha通道
                    color_rgba = (color[2], color[1], color[0], 255)
                    
                    # 绘制边界框
                    line_width = self.line_thickness
                    draw.rectangle([x1, y1, x2, y2], outline=(color_rgba[0], color_rgba[1], color_rgba[2], 255), width=line_width)
                    
                    # 绘制标签背景
                    text = f"{label} {confidence:.2f}"
                    
                    # 获取文本尺寸
                    if self.font:
                        try:
                            # 尝试使用textbbox方法（Pillow>=8.0.0）
                            text_bbox = draw.textbbox((0, 0), text, font=self.font)
                            text_width = text_bbox[2] - text_bbox[0]
                            text_height = text_bbox[3] - text_bbox[1]
                        except AttributeError:
                            # 兼容Pillow<8.0.0
                            text_width, text_height = draw.textsize(text, font=self.font)
                        
                        # 绘制标签背景
                        text_bg_x1 = x1
                        text_bg_y1 = y1 - text_height - 5
                        text_bg_x2 = x1 + text_width + 10
                        text_bg_y2 = y1
                        
                        # 如果标签超出屏幕顶部，将其放在边界框下方
                        if text_bg_y1 < 0:
                            text_bg_y1 = y2
                            text_bg_y2 = y2 + text_height + 5
                        
                        # 绘制半透明背景
                        draw.rectangle(
                            [text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2],
                            fill=(color_rgba[0], color_rgba[1], color_rgba[2], 180)
                        )
                        
                        # 计算文本位置（居中）
                        text_x = text_bg_x1 + 5
                        text_y = text_bg_y1 + 2
                        
                        # 绘制文本
                        draw.text((text_x, text_y), text, fill=(255, 255, 255, 255), font=self.font)
            
            # 添加FPS显示
            fps_text = f"FPS: {self.fps:.1f}"
            if self.font:
                draw.text((20, 20), fps_text, fill=(0, 255, 0, 255), font=self.font)
            
            # 将PIL图像转换为Tkinter图像并显示
            self.tk_image = ImageTk.PhotoImage(overlay)
            self.canvas.delete("all")
            self.canvas.create_image(0, 0, image=self.tk_image, anchor=tk.NW)
        except Exception as e:
            print(f"渲染检测结果时出错: {e}")
    
    def show_help(self):
        """显示帮助信息"""
        if not self.running or not self.canvas:
            return
        
        try:
            # 创建透明图像
            overlay = Image.new('RGBA', (self.screen_width, self.screen_height), (0, 0, 0, 0))
            draw = ImageDraw.Draw(overlay)
            
            # 帮助信息
            help_text = [
                "YOLOv8 检测覆盖层",
                "",
                "快捷键:",
                "  ESC 或 Q: 退出检测",
                "  +: 增加透明度",
                "  -: 减少透明度",
                "",
                "正在初始化，请稍候..."
            ]
            
            # 计算文本框位置（屏幕中央）
            text_height = len(help_text) * 30
            text_y = (self.screen_height - text_height) // 2
            
            # 绘制半透明背景
            bg_padding = 40
            bg_x1 = self.screen_width // 4
            bg_y1 = text_y - bg_padding
            bg_x2 = self.screen_width * 3 // 4
            bg_y2 = text_y + text_height + bg_padding
            
            # 绘制背景
            draw.rectangle(
                [bg_x1, bg_y1, bg_x2, bg_y2],
                fill=(0, 0, 0, 180),
                outline=(255, 255, 255, 255),
                width=2
            )
            
            # 绘制文本
            if self.font:
                for i, line in enumerate(help_text):
                    try:
                        text_width = draw.textbbox((0, 0), line, font=self.font)[2]
                    except AttributeError:
                        text_width, _ = draw.textsize(line, font=self.font)
                        
                    text_x = (self.screen_width - text_width) // 2
                    draw.text((text_x, text_y + i * 30), line, fill=(255, 255, 255, 255), font=self.font)
            
            # 将PIL图像转换为Tkinter图像并显示
            self.tk_image = ImageTk.PhotoImage(overlay)
            self.canvas.delete("all")
            self.canvas.create_image(0, 0, image=self.tk_image, anchor=tk.NW)
        except Exception as e:
            print(f"显示帮助信息时出错: {e}") 